"""
author：石沙
date：2020-09-28
content：本模块使用训练好的模型进行预测，供app.py调用
"""


from site_packages.utils.models import ModelOp
from .preprocess import BookInfoProcessor
from .features import ext_features


class Predictor:

    allowed_model_names = ['lgbm', 'bert', 'lr_voting']

    def __init__(self, model_name='lgbm'):
        self.model_name = model_name
        self.model = ModelOp.load_model(model_name)
        self.label_dict = ModelOp.load_data('label_dict')
        self.idx2label = dict(zip(self.label_dict.values(), self.label_dict.keys()))
        self.processor = BookInfoProcessor

    def get_features(self, inputs):
        inputs = self.processor(inputs, mode='answer').run()
        label_dict = ModelOp.load_data('label_dict')
        features = ext_features(inputs['full_content'].tolist(), label_dict, book_ids=None, return_df=True)
        print(features.shape)
        print(features.head())
        return features

    def predict(self, inputs):
        features = self.get_features(inputs)
        feature_cols = [col for col in features.columns if col not in ['book_id', 'label']]
        preds = self.model.predict(features[feature_cols].values)
        return {'label': self.idx2label[preds[0]]}

